package com.intmall.flink.state;

import com.intmall.flink.operator.ClickSource;
import com.intmall.flink.operator.Event;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.time.Duration;

public class PeriodicPvTest {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                        .withTimestampAssigner(new SerializableTimestampAssigner<Event>() {
                            @Override
                            public long extractTimestamp(Event element, long recordTimestamp) {
                                return element.getTimestamp();
                            }
                        }));

        stream.keyBy(Event::getUser)
                .process(new PeriodicPvResult(), Types.STRING)
                .print();
        env.execute();
    }

    public static class PeriodicPvResult extends KeyedProcessFunction<String, Event, String> {

        private ValueState<Long> countState;
        // 判断定时器是否存在，用于定时触发计算
        private ValueState<Long> timerTsState;

        @Override
        public void open(Configuration parameters) throws Exception {
            countState = getRuntimeContext().getState(new ValueStateDescriptor<Long>("count-state", Types.LONG));
            timerTsState = getRuntimeContext().getState(new ValueStateDescriptor<Long>("timer-ts", Types.LONG));
        }

        @Override
        public void processElement(Event value, KeyedProcessFunction<String, Event, String>.Context ctx, Collector<String> out) throws Exception {
            countState.update(countState.value() == null ? 1L : countState.value() + 1);

            // 如果没有注册定时器，就做注册
            if (timerTsState.value() == null) {
                ctx.timerService().registerEventTimeTimer(value.getTimestamp() + 10 * 1000L);
                timerTsState.update(value.getTimestamp() + 10 * 1000L);
            }
        }

        @Override
        public void onTimer(long timestamp, KeyedProcessFunction<String, Event, String>.OnTimerContext ctx, Collector<String> out) throws Exception {
            out.collect(ctx.getCurrentKey() + " pv count: " + countState.value());
            // 每次等下次的数据再来时再注册定时器(中间等待下次数据来时会有间隔)
            timerTsState.clear();

            // 如果需要固定时间触发，在这里重新注册定时器
//            ctx.timerService().registerEventTimeTimer(timestamp + 10 * 1000L);
        }
    }
}
